import queue
import av
import cv2
from yolo_core import YOLODetector
from rfdetr_core import RFDETRDetector

# 初始化检测器，使用TensorRT模型
detector = YOLODetector('models/best.engine')
rfdetr_detector = RFDETRDetector()

# 目标检测处理函数
def yolo_frame(rtc_queue, yolo_queue):
    while True:
        try:
            # 打印两个队列的长度
            print(f"原始队列长度: {rtc_queue.qsize()}, 检测队列长度: {yolo_queue.qsize()}")

            # 从原始视频帧队列中获取帧
            frame = rtc_queue.get()

            #进行中文检测
            # annotated_frame = detector.detect_and_draw_Chinese(frame)

            #进行英文检测
            # annotated_frame = detector.detect_and_draw_English(frame)

            # 使用带计数功能的检测方法
            annotated_frame = rfdetr_detector.detect_and_draw_count(frame)
            # 将检测结果写入到检测结果队列中
            yolo_queue.put(annotated_frame)
        except Exception as e:
            print(f"Error in yolo_frame: {e}")


# 用于从检测队列中获取帧并推送到 RTMP 服务器
def push_frame(yolo_queue, rtmp_url):
    # 先获取第一个帧以确定宽高
    first_frame = yolo_queue.get()
    height, width, _ = first_frame.shape

    # 创建输出容器
    output_container = av.open(rtmp_url, 'w', format='flv')
    # 创建视频流
    stream = output_container.add_stream('libx264', rate=25)
    stream.pix_fmt = 'yuv420p'
    stream.width = width
    stream.height = height
    stream.options = {'preset': 'ultrafast', 'tune': 'zerolatency', 'crf': '25'}
    # stream.codec_context.options = {'tune': 'zerolatency', 'x264-params': 'no-scenecut'}

    try:
        # 先处理第一个帧
        video_frame = av.VideoFrame.from_ndarray(first_frame, format='bgr24')
        for packet in stream.encode(video_frame):
            output_container.mux(packet)

        while True:
            print(f"推送时的队列长度: {yolo_queue.qsize()}")
            # 从检测队列中获取一个帧【格式是 BGR24 格式的 NumPy 数组】
            frame = yolo_queue.get()
            # 将 NumPy 数组转换为 av.VideoFrame
            video_frame = av.VideoFrame.from_ndarray(frame, format='bgr24')
            # 将视频帧编码
            for packet in stream.encode(video_frame):
                output_container.mux(packet)
    except Exception as e:
        print(f"Error in push_frame: {e}")
    finally:
        # 清理资源
        for packet in stream.encode(None):
            output_container.mux(packet)
        output_container.close()
